import cv2
import numpy as np
from skimage import filters


def getOTSUMask(data_array):
    data_array_copy = data_array.astype(np.uint8)
    mask = np.zeros_like(data_array_copy)
    for i in range(data_array_copy.shape[0]):
        tmp = cv2.threshold(data_array_copy[i, :, :], 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        tmp = tmp[1]  # 0 is the threshold
        mask[i, :, :] = tmp // 255
    # data_array = np.multiply(data_array, mask)
    return mask


def myAdaptiveThreshold(data_array):
    data_array = data_array * 255
    data_array = data_array.astype(np.uint8)
    result = np.zeros_like(data_array)
    for i in range(data_array.shape[0]):
        tmp = cv2.adaptiveThreshold(data_array[i, :, :], 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 7, 0)
        result[i, :, :] = tmp // 255
    return result


def simple_threshold(data_array, threshold):
    return np.where(data_array > threshold, 1, 0)


def hysteresis_threshold(predict, low_threshold=0.4, high_threshold=0.5):
    result = np.zeros_like(predict)
    if len(predict.shape) == 5:
        tmp_thresholded = filters.apply_hysteresis_threshold(predict[0, 0, :, :, :], low_threshold, high_threshold)
        result[0, 0, :, :, :] = tmp_thresholded
    elif len(predict.shape) == 3:
        result = filters.apply_hysteresis_threshold(predict, low_threshold, high_threshold)
    # The return of function apply_hysteresis_threshold is bool
    return result.astype(np.long)
